import cv2
import numpy as np

'传入预先利用棋盘格检测出来的相机内参'
'问的AI，一开始强行使用猜测的值然后鲁棒性不太行，换了这个棋盘格算法'
def load_camera_params(path="camera_calib.npz"):
    data = np.load(path)
    return data['cameramatrix'], data['distcoeffs']

'四个角点的排序，左上 右上 右下 左下,生成坐标'
def sort_corners(pts):
    pts = pts.reshape(4, 2)
    s = pts.sum(axis=1)
    diff = np.diff(pts, axis=1)
    ordered = np.zeros((4, 2), dtype=np.float32)
    ordered[0] = pts[np.argmin(s)]  # 左上 x+y=min
    ordered[2] = pts[np.argmax(s)]  # 右下 x+y=max
    ordered[1] = pts[np.argmin(diff)]  # 右上 x-y=min
    ordered[3] = pts[np.argmax(diff)]  # 左下 x-y=max
    return ordered

def detect_rectangle(gray):
    blur = cv2.GaussianBlur(gray, (5, 5), 0)#去噪
    '查资料改用canny检测，提高鲁棒性'
    edges = cv2.Canny(blur, 50, 150)
    kernel = np.ones((3, 3), np.uint8)
    edges = cv2.dilate(edges, kernel, iterations=1)
    '算法同最小多边形拟合'
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    max_area, best = 0, None
    for cnt in contours:
        area = cv2.contourArea(cnt)
        peri = cv2.arcLength(cnt, True)
        '多边形拟合中的近似'
        approx = cv2.approxPolyDP(cnt, 0.02 * peri, True)
        '只要凸四边形'
        if len(approx) == 4 and cv2.isContourConvex(approx):
            '修改阈值到2000，提高可靠度'
            if area < 2000:
                continue
            if area > max_area:
                best = approx
                max_area = area
    return sort_corners(best) \
        if best is not None else None#只要最大的矩形，提高可靠度

'算距离'
def distance(p1, p2):
    return np.linalg.norm(np.array(p1) - np.array(p2))

'3D图形中的坐标定义'
def draw_axes(img, cameramatrix, distcoeffs, rvec, tvec, long_len, short_len):
    axis = np.float32([
        [0, 0, 0],
        [long_len, 0, 0],
        [0, -short_len, 0],
        [0, 0, max(long_len, short_len) / 2]
    ])
    '2D中的投影'
    imgpts, _ = cv2.projectPoints(axis, rvec, tvec, cameramatrix, distcoeffs)
    if imgpts is None or imgpts.shape != (4,1,2) or not np.isfinite(imgpts).all():
        return
    imgpts = imgpts.reshape(-1, 2)
    try:
        o = tuple(imgpts[0].astype(int))
        x = tuple(imgpts[1].astype(int))
        y = tuple(imgpts[2].astype(int))
        z = tuple(imgpts[3].astype(int))
        cv2.line(img, o, x, (0, 0, 255), 3)#红
        cv2.line(img, o, y, (0, 255, 0), 3)#绿
        cv2.line(img, o, z, (255, 0, 0), 3)#蓝
    except:
        pass


def main():
    cam_mtx, dist = load_camera_params()
    cap = cv2.VideoCapture(0)
    '矩形实际尺寸'
    long_len, short_len = 170.0, 80.0

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        corners = detect_rectangle(gray)
        '框选矩形出来，看看到底是PnP问题还是内参问题'
        if corners is not None:
            cv2.polylines(frame, [corners.reshape(-1, 1, 2).astype(int)], True, (200, 200, 200), 2)
            tl, tr, br, bl = corners
            len_top = distance(tl, tr)
            len_left = distance(tl, bl)
            long_is_horizontal = len_top > len_left

            '找长边是否水平'
            if long_is_horizontal:
                obj_pts = np.array([
                    [0, 0, 0],
                    [long_len, 0, 0],
                    [long_len, short_len, 0],
                    [0, short_len, 0]
                ], dtype=np.float32)
                img_pts = np.array([tl, tr, br, bl], dtype=np.float32)
            else:
                obj_pts = np.array([
                    [0, 0, 0],
                    [long_len, 0, 0],
                    [long_len, short_len, 0],
                    [0, short_len, 0]
                ], dtype=np.float32)
                img_pts = np.array([tl, bl, br, tr], dtype=np.float32)

            'PnP算法'
            success, rvec, tvec = cv2.solvePnP(obj_pts, img_pts, cam_mtx, dist, flags=cv2.SOLVEPNP_IPPE)
            if success:
                r, _ = cv2.Rodrigues(rvec)
                z_cam = r[:, 2]
                '确保Z轴能正常显示而不是和初版一样乱飞'
                if z_cam[2] > 0:
                    r = r @ np.diag([1, -1, -1])
                    rvec, _ = cv2.Rodrigues(r)
                draw_axes(frame, cam_mtx, dist, rvec, tvec, long_len, short_len)

        cv2.imshow("Pose", frame)
        if cv2.waitKey(1) == 27:
            break#ESC的ASCII码27

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
